import os
import numpy as np
from PIL import Image
from sklearn.model_selection import train_test_split
from sklearn.svm import SVC
import joblib

# 加载数据
data_dir = 'captcha_images'
captcha_images = []
captcha_labels = []
for filename in os.listdir(data_dir):
    if filename.endswith('.png'):
        captcha_text = os.path.splitext(filename)[0]
        captcha_image = Image.open(os.path.join(data_dir, filename)).convert('L')
        captcha_image = np.array(captcha_image).flatten()
        captcha_images.append(captcha_image)
        captcha_labels.append(captcha_text)
captcha_images = np.array(captcha_images)
captcha_labels = np.array(captcha_labels)

# 划分数据集
x_train, x_test, y_train, y_test = train_test_split(captcha_images, captcha_labels, test_size=0.2)

# 构建模型
model = SVC(kernel='rbf', C=1, gamma='scale')

# 训练模型
model.fit(x_train, y_train)

# 评估模型
train_score = model.score(x_train, y_train)
test_score = model.score(x_test, y_test)
print('Train score:', train_score)
print('Test score:', test_score)

# 保存模型
joblib.dump(model, 'captcha_model.pkl')